"""
MCP Server 实现
"""
import asyncio
import json
import os
import sys
from typing import Any, Dict, List, Optional

from mcp.server.fastmcp import FastMCP
import uuid

# 添加项目路径到sys.path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.database import SessionLocal
from models.database import User, ParseTask, SliceTask
from services.parse_service import ParseService
from services.slice_service import SliceService
from core.auth import generate_task_id, generate_slice_id
from core.config import get_settings

settings = get_settings()

# 创建FastMCP服务器实例
mcp = FastMCP(
    "MinerU-Doc-Parser", 
    port=settings.mcp_port,
    host=settings.host,
    debug=settings.debug
)


class MinerUMCPServer:
    """MinerU MCP服务器"""
    
    def __init__(self):
        self.parse_service = ParseService()
        self.slice_service = SliceService()


# 创建MCP服务器实例
mcp_server = MinerUMCPServer()


@mcp.tool()
async def parse_document(
    file_path: str, 
    options: Optional[Dict[str, Any]] = None
) -> str:
    """解析文档工具
    
    Args:
        file_path: 要解析的文档文件路径
        options: 解析选项，包括输出格式、语言、OCR引擎等
        
    Returns:
        JSON格式的解析结果
    """
    try:
        # 验证文件路径
        if not os.path.exists(file_path):
            return json.dumps({"error": f"File not found: {file_path}"}, ensure_ascii=False)
        
        # 生成任务ID
        task_id = generate_task_id()
        
        # 默认选项
        default_options = {
            "output_format": "all",
            "language": "auto",
            "ocr_engine": "paddle"
        }
        
        if options:
            default_options.update(options)
        
        # 创建数据库会话
        db = SessionLocal()
        
        try:
            # 创建默认用户（如果不存在）
            user = db.query(User).filter(User.username == "mcp_user").first()
            if not user:
                user = User(
                    username="mcp_user",
                    email="mcp@system.com",
                    password_hash="mcp_password_hash",
                    role="user"
                )
                db.add(user)
                db.commit()
                db.refresh(user)
            
            # 创建解析任务
            parse_task = ParseTask(
                task_id=task_id,
                user_id=user.id,
                filename=os.path.basename(file_path),
                file_path=file_path,
                file_size=os.path.getsize(file_path),
                status="processing",
                execution_mode="sync",
                parse_options=default_options
            )
            
            db.add(parse_task)
            db.commit()
            db.refresh(parse_task)
            
            # 执行解析
            from core.logger import get_task_logger
            task_logger = get_task_logger(task_id)
            
            result = await mcp_server.parse_service.parse_document(
                file_path=file_path,
                task_id=task_id,
                options=default_options,
                task_logger=task_logger
            )
            
            # 更新任务状态
            parse_task.status = "completed"
            parse_task.progress = 100
            parse_task.result_path = result.get("result_path")
            db.commit()
            
            # 保存解析结果
            from models.database import ParseResult
            for content_type, content_info in result.get("results", {}).items():
                parse_result = ParseResult(
                    task_id=parse_task.id,
                    content_type=content_type,
                    content_path=content_info.get("path"),
                    file_size=content_info.get("size"),
                    metadata=content_info.get("metadata")
                )
                db.add(parse_result)
            
            db.commit()
            
            return json.dumps({
                "success": True,
                "task_id": task_id,
                "results": result.get("results", {}),
                "download_urls": result.get("download_urls", {})
            }, ensure_ascii=False)
            
        except Exception as e:
            # 更新任务状态为失败
            if 'parse_task' in locals():
                parse_task.status = "failed"
                parse_task.error_message = str(e)
                db.commit()
            
            return json.dumps({"error": f"Parse failed: {str(e)}"}, ensure_ascii=False)
        finally:
            db.close()
            
    except Exception as e:
        return json.dumps({"error": f"Tool execution failed: {str(e)}"}, ensure_ascii=False)


@mcp.tool()
async def slice_document(
    task_id: str, 
    options: Optional[Dict[str, Any]] = None
) -> str:
    """文档切片工具
    
    Args:
        task_id: 已完成的解析任务ID
        options: 切片选项，包括策略、块大小、重叠等
        
    Returns:
        JSON格式的切片结果
    """
    try:
        # 默认选项
        default_options = {
            "strategy": "semantic",
            "chunk_size": 1000,
            "overlap": 200,
            "min_chunk_size": 100
        }
        
        if options:
            default_options.update(options)
        
        # 创建数据库会话
        db = SessionLocal()
        
        try:
            # 获取解析任务
            parse_task = db.query(ParseTask).filter(ParseTask.task_id == task_id).first()
            if not parse_task:
                return json.dumps({"error": f"Parse task not found: {task_id}"}, ensure_ascii=False)
            
            if parse_task.status != "completed":
                return json.dumps({"error": f"Parse task not completed: {task_id}"}, ensure_ascii=False)
            
            # 生成切片ID
            slice_id = generate_slice_id()
            
            # 创建切片任务
            slice_task = SliceTask(
                slice_id=slice_id,
                parse_task_id=parse_task.id,
                status="processing",
                strategy=default_options["strategy"],
                chunk_size=default_options["chunk_size"],
                overlap=default_options["overlap"],
                min_chunk_size=default_options["min_chunk_size"]
            )
            
            db.add(slice_task)
            db.commit()
            db.refresh(slice_task)
            
            # 执行切片
            from core.logger import get_task_logger
            task_logger = get_task_logger(slice_id)
            
            result = await mcp_server.slice_service.slice_document(
                parse_task_id=parse_task.id,
                options=default_options,
                slice_id=slice_id,
                task_logger=task_logger,
                db=db
            )
            
            # 更新任务状态
            slice_task.status = "completed"
            slice_task.total_chunks = len(result.get("chunks", []))
            slice_task.result_path = result.get("result_path")
            db.commit()
            
            # 保存切片块
            from models.database import SliceChunk
            for chunk_data in result.get("chunks", []):
                chunk = SliceChunk(
                    slice_task_id=slice_task.id,
                    chunk_id=chunk_data.get("chunk_id"),
                    text_content=chunk_data.get("text_content"),
                    word_count=chunk_data.get("word_count", 0),
                    page_number=chunk_data.get("page_number"),
                    position_info=chunk_data.get("position_info"),
                    metadata=chunk_data.get("metadata")
                )
                db.add(chunk)
            
            db.commit()
            
            return json.dumps({
                "success": True,
                "slice_id": slice_id,
                "chunks": result.get("chunks", []),
                "total_chunks": len(result.get("chunks", [])),
                "statistics": await mcp_server.slice_service.get_slice_statistics(slice_id, db)
            }, ensure_ascii=False)
            
        except Exception as e:
            # 更新任务状态为失败
            if 'slice_task' in locals():
                slice_task.status = "failed"
                slice_task.error_message = str(e)
                db.commit()
            
            return json.dumps({"error": f"Slice failed: {str(e)}"}, ensure_ascii=False)
        finally:
            db.close()
            
    except Exception as e:
        return json.dumps({"error": f"Tool execution failed: {str(e)}"}, ensure_ascii=False)


@mcp.tool()
async def get_task_status(task_id: str) -> str:
    """获取任务状态工具
    
    Args:
        task_id: 任务ID（解析任务或切片任务）
        
    Returns:
        JSON格式的任务状态信息
    """
    try:
        db = SessionLocal()
        
        try:
            # 查找解析任务
            parse_task = db.query(ParseTask).filter(ParseTask.task_id == task_id).first()
            if parse_task:
                result = {
                    "type": "parse",
                    "task_id": task_id,
                    "status": parse_task.status,
                    "progress": parse_task.progress,
                    "filename": parse_task.filename,
                    "created_at": parse_task.created_at.isoformat(),
                    "started_at": parse_task.started_at.isoformat() if parse_task.started_at else None,
                    "completed_at": parse_task.completed_at.isoformat() if parse_task.completed_at else None,
                    "error_message": parse_task.error_message,
                    "execution_mode": parse_task.execution_mode
                }
                return json.dumps(result, ensure_ascii=False)
            
            # 查找切片任务
            slice_task = db.query(SliceTask).filter(SliceTask.slice_id == task_id).first()
            if slice_task:
                result = {
                    "type": "slice",
                    "task_id": task_id,
                    "status": slice_task.status,
                    "strategy": slice_task.strategy,
                    "chunk_size": slice_task.chunk_size,
                    "overlap": slice_task.overlap,
                    "total_chunks": slice_task.total_chunks,
                    "created_at": slice_task.created_at.isoformat(),
                    "completed_at": slice_task.completed_at.isoformat() if slice_task.completed_at else None,
                    "error_message": slice_task.error_message
                }
                return json.dumps(result, ensure_ascii=False)
            
            return json.dumps({"error": f"Task not found: {task_id}"}, ensure_ascii=False)
            
        finally:
            db.close()
            
    except Exception as e:
        return json.dumps({"error": f"Tool execution failed: {str(e)}"}, ensure_ascii=False)


@mcp.tool()
async def get_parse_results(
    task_id: str, 
    format: Optional[str] = None
) -> str:
    """获取解析结果工具
    
    Args:
        task_id: 解析任务ID
        format: 可选，指定返回的内容格式
        
    Returns:
        JSON格式的解析结果
    """
    try:
        db = SessionLocal()
        
        try:
            # 获取解析任务
            parse_task = db.query(ParseTask).filter(ParseTask.task_id == task_id).first()
            if not parse_task:
                return json.dumps({"error": f"Parse task not found: {task_id}"}, ensure_ascii=False)
            
            if parse_task.status != "completed":
                return json.dumps({"error": f"Parse task not completed: {task_id}"}, ensure_ascii=False)
            
            # 获取解析结果
            from models.database import ParseResult
            results = db.query(ParseResult).filter(ParseResult.task_id == parse_task.id).all()
            
            if not results:
                return json.dumps({"error": f"No results found for task: {task_id}"}, ensure_ascii=False)
            
            # 根据格式返回结果
            if format:
                result = next((r for r in results if r.content_type == format), None)
                if not result:
                    return json.dumps({"error": f"Result format '{format}' not found"}, ensure_ascii=False)
                
                # 读取文件内容
                if os.path.exists(result.content_path):
                    with open(result.content_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    return json.dumps({
                        "content_type": result.content_type,
                        "content": content,
                        "metadata": result.result_metadata,
                        "file_size": result.file_size
                    }, ensure_ascii=False)
                else:
                    return json.dumps({"error": f"Result file not found: {result.content_path}"}, ensure_ascii=False)
            else:
                # 返回所有结果的信息
                result_info = []
                for result in results:
                    info = {
                        "content_type": result.content_type,
                        "file_size": result.file_size,
                        "metadata": result.result_metadata,
                        "created_at": result.created_at.isoformat()
                    }
                    result_info.append(info)
                
                return json.dumps({
                    "task_id": task_id,
                    "results": result_info,
                    "total_results": len(result_info)
                }, ensure_ascii=False)
            
        finally:
            db.close()
            
    except Exception as e:
        return json.dumps({"error": f"Tool execution failed: {str(e)}"}, ensure_ascii=False)


@mcp.tool()
async def get_slice_results(
    slice_id: str, 
    page: Optional[int] = None, 
    limit: Optional[int] = None
) -> str:
    """获取切片结果工具
    
    Args:
        slice_id: 切片任务ID
        page: 可选，页码（用于分页）
        limit: 可选，每页数量
        
    Returns:
        JSON格式的切片结果
    """
    try:
        db = SessionLocal()
        
        try:
            # 获取切片任务
            slice_task = db.query(SliceTask).filter(SliceTask.slice_id == slice_id).first()
            if not slice_task:
                return json.dumps({"error": f"Slice task not found: {slice_id}"}, ensure_ascii=False)
            
            if slice_task.status != "completed":
                return json.dumps({"error": f"Slice task not completed: {slice_id}"}, ensure_ascii=False)
            
            # 获取切片块
            from models.database import SliceChunk
            query = db.query(SliceChunk).filter(SliceChunk.slice_task_id == slice_task.id)
            
            # 分页
            if page and limit:
                offset = (page - 1) * limit
                query = query.offset(offset).limit(limit)
            
            chunks = query.all()
            
            # 转换为字典格式
            chunk_data = []
            for chunk in chunks:
                chunk_dict = {
                    "chunk_id": chunk.chunk_id,
                    "text_content": chunk.text_content,
                    "word_count": chunk.word_count,
                    "page_number": chunk.page_number,
                    "position_info": chunk.position_info,
                    "metadata": chunk.chunk_metadata,
                    "created_at": chunk.created_at.isoformat()
                }
                chunk_data.append(chunk_dict)
            
            # 获取统计信息
            statistics = await mcp_server.slice_service.get_slice_statistics(slice_id, db)
            
            return json.dumps({
                "slice_id": slice_id,
                "chunks": chunk_data,
                "total_chunks": len(chunk_data),
                "page": page,
                "limit": limit,
                "statistics": statistics
            }, ensure_ascii=False)
            
        finally:
            db.close()
            
    except Exception as e:
        return json.dumps({"error": f"Tool execution failed: {str(e)}"}, ensure_ascii=False)


if __name__ == "__main__":
    # 初始化数据库
    from core.database import init_db
    init_db()
    
    # 启动MCP服务器
    print(f"启动MCP HTTP服务器，端口: {settings.mcp_port}")
    mcp.run(transport="streamable-http")